import os
import sys
import time
import toml
from mininet.topo import Topo
from mininet.net import Mininet
from mininet.util import dumpNodeConnections
from mininet.log import setLogLevel
from mininet.util import irange
from mininet.cli import CLI
from mininet.link import TCLink
from typing import Dict, List, Union, Any


port_start = 9000
in_path = 'temp'

class MyTopo(Topo):
    def __init__(self, config):
        self._config = config
        self._flows = []
        self._ports = {}
        super().__init__()

    def build(self, **_opts):
        num_hosts = self._config['num_hosts']
        num_switches = self._config['num_switches']
        # Add hosts and switches
        for i in range(num_hosts):
            self.addHost(f'h{i}')
        for i in range(num_switches):
            self.addSwitch(f's{i}') 

        default_config = [i for i in self._config['link_properties'] if i['name'] == 'default']
        if default_config:
            default_config = default_config[0].copy()
            del default_config['name']
        
        # Add links
        self._ports = {f's{i}': {} for i in range(num_switches)}
        ports = self._ports

        for i in self._config['topo']:
            if 'property' in i:
                name = i['property']
                config = [j for j in self._config['link_properties'] if j['name'] == name]
                if config:
                    config = config[0].copy()
                    del config['name']
                else:
                    raise RuntimeError('Config name does not exist.')
            elif default_config:
                config = default_config
            else:
                raise RuntimeError('You need to specify a config for a flow.')
            node1, node2 = i['link']
            self.addLink(node1, node2, **config)
            if node1[0] == 's':
                ports[node1][node2] = len(ports[node1]) + 1
            if node2[0] == 's':
                ports[node2][node1] = len(ports[node2]) + 1

        # Set flow rules
        

    def get_flows(self):
        return self._flows
    
    def get_ports(self):
        return self._ports


def set_rules(net: Mininet):
    topo: MyTopo = net.topo
    src_dst = {}
    for i in topo._config['flows']:
        path = i['path']
        pair = (path[0], path[-1])
        if pair not in src_dst:
            topo.get_flows().append({**i, 'index': 0})
            src_dst[pair] = 1
        else:
            topo.get_flows().append({**i, 'index': src_dst[pair]})
            src_dst[pair] += 1
    for i in topo.get_flows():
        src_name = i['path'][0]
        src_ip = net.getNodeByName(src_name).IP()
        dst_name = i['path'][-1]
        dst_ip = net.getNodeByName(dst_name).IP()
        tcp_port = port_start + i['index']
        for j in range(1, len(i['path']) - 1):
            switch_name = i['path'][j]
            switch = net.getNodeByName(switch_name)
            next_node = i['path'][j+1]
            prev_node = i['path'][j-1]
            out_port_next = topo.get_ports()[switch_name][next_node]
            out_port_prev = topo.get_ports()[switch_name][prev_node]
            switch.cmd(f'ovs-ofctl add-flow {switch_name} dl_type:0x800,nw_proto=6,nw_src:{src_ip},nw_dst:{dst_ip},tp_dst:{tcp_port},actions=output:{out_port_next}')
            switch.cmd(f'ovs-ofctl add-flow {switch_name} dl_type:0x800,nw_proto=6,nw_src:{dst_ip},nw_dst:{src_ip},tp_src:{tcp_port},actions=output:{out_port_prev}')

def start_test(net: Mininet):
    topo: MyTopo = net.topo
    servers = set()
    for i, f in enumerate(topo.get_flows()):
        src_name = f['path'][0]
        src = net.getNodeByName(src_name)
        dst_name = f['path'][-1]
        dst = net.getNodeByName(dst_name)
        dst_ip = net.getNodeByName(dst_name).IP()
        tcp_port = port_start + f['index']
        if (dst, tcp_port) not in servers:
            servers.add((dst, tcp_port))
            dst.cmd(f'iperf -s -p {tcp_port} &')
        src.cmd(f'iperf -c {dst_ip} -p {tcp_port} -t 10 -i 1 > {in_path}/flow{i}.log &')

def output(in_path: str, out_path: str, num_flows: int):
    bandwidth = ['' for i in range(num_flows)]
    for root, _, files in os.walk(in_path):
        for file in files:
            i = int(file.split('.')[0][4:])
            with open(os.path.join(root, file)) as f:
                content = f.readlines()[-1]
            bw = content.split(' ')[-2]
            bandwidth[i] = bw
    with open(out_path, 'w') as f:
        for i, bw in enumerate(bandwidth):
            f.write(f'{i}\t{bw}\n')


if __name__ == '__main__':
    if not os.path.exists(in_path):
        os.mkdir(in_path)

    config = toml.load(sys.argv[1])
    print(config)

    setLogLevel('info')
    topo = MyTopo(config)
    net = Mininet(topo=topo, link=TCLink, autoStaticArp=True)
    net.start()
    set_rules(net)
    start_test(net)
    print('start simulation')
    time.sleep(15)
    output(in_path, sys.argv[2], len(topo._config['flows']))
    net.stop()

